package first.group.ssh.keyexchange;

import first.group.ssh.packet.writer.BigIntegerWriter;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.security.DigestException;
import java.security.MessageDigest;

public class KeyGenerator {

    public static class Entry {

        private final MessageDigest digest;
        private final byte[] pack; // K || H
        private final byte ch;
        private final byte[] sessionId; // session_id

        private Entry(MessageDigest digest, byte[] pack, byte ch, byte[] sessionId) {
            this.digest = digest;
            this.pack = pack;
            this.ch = ch;
            this.sessionId = sessionId;
        }

        public byte[] generate(int len) {
            try {
                digest.reset();
                digest.update(pack);
                digest.update(ch);
                digest.update(sessionId);

                var digestLen = digest.getDigestLength();
                if (digestLen == 0) {
                    throw new IllegalArgumentException("Cannot get digest length");
                }

                var count = len / digestLen;
                if (len % digestLen != 0) {
                    count++;
                }
                var buffer = new byte[count * digestLen];

                var leftLen = len;
                var offset = 0;
                while (leftLen > 0) {
                    var outSize = Math.min(digestLen, leftLen);
                    digest.digest(buffer, offset, digestLen);
                    leftLen -= outSize;
                    if (leftLen != 0) {
                        digest.reset();
                        digest.update(pack);
                        digest.update(buffer, offset - digestLen, digestLen);
                    }
                }

                var result = new byte[len];
                System.arraycopy(buffer, 0, result, 0, len);
                return result;
            } catch (DigestException ex) {
                throw new RuntimeException(ex);
            }
        }
    }

    public static byte[] pack(BigInteger k, byte[] h) {
        var bitLen = k.bitLength() + 1;
        var byteLen = bitLen / 8;
        if (bitLen % 8 != 0) {
            byteLen += 1;
        }

        var buffer = ByteBuffer.allocate(byteLen + h.length + 4);
        buffer.order(ByteOrder.BIG_ENDIAN);

        BigIntegerWriter.writeBigInteger(buffer, k);
        buffer.put(h);

        return buffer.array();
    }

    public static Entry newEntry(KeyExchangeMethod method, byte[] pack, byte ch, byte[] sessionId) {
        return new Entry(method.newDigest(), pack, ch, sessionId);
    }

    public static byte[] generate(KeyExchangeMethod method, BigInteger k, byte[] h, byte ch, byte[] sessionId, int len) {
        var entry = newEntry(method, pack(k, h), ch, sessionId);
        return entry.generate(len);
    }

    public static byte[] generate(KeyExchangeMethod method, byte[] pack, byte ch, byte[] sessionId, int len) {
        var entry = newEntry(method, pack, ch, sessionId);
        return entry.generate(len);
    }
}
